/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.handler.component;

import static org.apache.solr.common.params.CommonParams.FQ;
import static org.apache.solr.common.params.CommonParams.JSON;

import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;
import org.apache.lucene.search.Query;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.util.SimpleOrderedMap;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.search.DocList;
import org.apache.solr.search.QueryParsing;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.search.facet.FacetDebugInfo;
import org.apache.solr.search.stats.StatsCache;
import org.apache.solr.util.SolrPluginUtils;
import org.apache.solr.util.SolrResponseUtil;

/**
 * Adds debugging information to a request.
 *
 * @since solr 1.3
 */
public class DebugComponent extends SearchComponent {
  public static final String COMPONENT_NAME = "debug";

  /**
   * Map containing all the possible stages as key and the corresponding readable purpose as value
   */
  private static final Map<Integer, String> stages;

  static {
    Map<Integer, String> map = new TreeMap<>();
    map.put(ResponseBuilder.STAGE_START, "START");
    map.put(ResponseBuilder.STAGE_PARSE_QUERY, "PARSE_QUERY");
    map.put(ResponseBuilder.STAGE_TOP_GROUPS, "TOP_GROUPS");
    map.put(ResponseBuilder.STAGE_EXECUTE_QUERY, "EXECUTE_QUERY");
    map.put(ResponseBuilder.STAGE_GET_FIELDS, "GET_FIELDS");
    map.put(ResponseBuilder.STAGE_DONE, "DONE");
    stages = Collections.unmodifiableMap(map);
  }

  @Override
  public void prepare(ResponseBuilder rb) throws IOException {
    if (rb.isDebugTrack() && rb.isDistrib) {
      rb.setNeedDocList(true);
      doDebugTrack(rb);
    }
  }

  @Override
  public void process(ResponseBuilder rb) throws IOException {
    if (rb.isDebug()) {
      SolrQueryRequest req = rb.req;
      StatsCache statsCache = req.getSearcher().getStatsCache();
      req.getContext().put(SolrIndexSearcher.STATS_SOURCE, statsCache.get(req));
      DocList results = null;
      // some internal grouping requests won't have results value set
      if (rb.getResults() != null) {
        results = rb.getResults().docList;
      }

      NamedList<Object> stdinfo =
          SolrPluginUtils.doStandardDebug(
              rb.req,
              rb.getQueryString(),
              rb.wrap(rb.getQuery()),
              results,
              rb.isDebugQuery(),
              rb.isDebugResults());

      NamedList<Object> info = rb.getDebugInfo();
      if (info == null) {
        rb.setDebugInfo(stdinfo);
        info = stdinfo;
      } else {
        info.addAll(stdinfo);
      }

      FacetDebugInfo fdebug = (FacetDebugInfo) (rb.req.getContext().get("FacetDebugInfo"));
      if (fdebug != null) {
        info.add("facet-trace", fdebug.getFacetDebugInfo());
      }

      fdebug = (FacetDebugInfo) (rb.req.getContext().get("FacetDebugInfo-nonJson"));
      if (fdebug != null) {
        info.add("facet-debug", fdebug.getFacetDebugInfo());
      }

      if (rb.req.getJSON() != null) {
        info.add(JSON, rb.req.getJSON());
      }

      if (rb.isDebugQuery() && rb.getQparser() != null) {
        rb.getQparser().addDebugInfo(rb.getDebugInfo());
      }

      if (null != rb.getDebugInfo()) {
        if (rb.isDebugQuery() && null != rb.getFilters()) {
          info.add("filter_queries", rb.req.getParams().getParams(FQ));
          List<String> fqs = new ArrayList<>(rb.getFilters().size());
          for (Query fq : rb.getFilters()) {
            fqs.add(QueryParsing.toString(fq, rb.req.getSchema()));
          }
          info.add("parsed_filter_queries", fqs);
        }

        // Add this directly here?
        rb.rsp.add("debug", rb.getDebugInfo());
      }
    }
  }

  private void doDebugTrack(ResponseBuilder rb) {
    final String rid = rb.req.getParams().get(CommonParams.REQUEST_ID);
    rb.addDebug(rid, "track", CommonParams.REQUEST_ID); // to see it in the response
  }

  @Override
  public void modifyRequest(ResponseBuilder rb, SearchComponent who, ShardRequest sreq) {
    if (!rb.isDebug()) return;

    // Turn on debug to get explain only when retrieving fields
    if ((sreq.purpose & ShardRequest.PURPOSE_GET_FIELDS) != 0) {
      sreq.purpose |= ShardRequest.PURPOSE_GET_DEBUG;
      // always distribute the latest version of global stats
      sreq.purpose |= ShardRequest.PURPOSE_SET_TERM_STATS;
      StatsCache statsCache = rb.req.getSearcher().getStatsCache();
      statsCache.sendGlobalStats(rb, sreq);

      if (rb.isDebugAll()) {
        sreq.params.set(CommonParams.DEBUG_QUERY, "true");
      } else {
        if (rb.isDebugQuery()) {
          sreq.params.add(CommonParams.DEBUG, CommonParams.QUERY);
        }
        if (rb.isDebugResults()) {
          sreq.params.add(CommonParams.DEBUG, CommonParams.RESULTS);
        }
      }
    } else {
      sreq.params.set(CommonParams.DEBUG_QUERY, "false");
      sreq.params.set(CommonParams.DEBUG, "false");
    }
    if (rb.isDebugTimings()) {
      sreq.params.add(CommonParams.DEBUG, CommonParams.TIMING);
    }
    if (rb.isDebugTrack()) {
      sreq.params.add(CommonParams.DEBUG, CommonParams.TRACK);
      sreq.params.set(CommonParams.REQUEST_ID, rb.req.getParams().get(CommonParams.REQUEST_ID));
      sreq.params.set(
          CommonParams.REQUEST_PURPOSE, SolrPluginUtils.getRequestPurpose(sreq.purpose));
    }
  }

  @VisibleForTesting
  protected String getDistributedStageName(int stage) {
    String stageName = stages.get(stage);

    if (stageName == null) {
      stageName = "STAGE_" + Integer.toString(stage);
    }

    return stageName;
  }

  @Override
  public void handleResponses(ResponseBuilder rb, ShardRequest sreq) {
    if (rb.isDebugTrack() && rb.isDistrib && !rb.finished.isEmpty()) {
      @SuppressWarnings("unchecked")
      NamedList<Object> stageList =
          (NamedList<Object>)
              ((NamedList<Object>) rb.getDebugInfo().get("track"))
                  .get(getDistributedStageName(rb.getStage()));
      if (stageList == null) {
        stageList = new SimpleOrderedMap<>();
        rb.addDebug(stageList, "track", getDistributedStageName(rb.getStage()));
      }
      for (ShardResponse response : sreq.responses) {
        stageList.add(response.getShard(), getTrackResponse(response));
      }
    }
  }

  private static final Set<String> EXCLUDE_SET = Set.of("explain");

  @Override
  @SuppressWarnings({"unchecked"})
  public void finishStage(ResponseBuilder rb) {
    if (rb.isDebug() && rb.getStage() == ResponseBuilder.STAGE_GET_FIELDS) {
      NamedList<Object> info = rb.getDebugInfo();
      NamedList<Object> explain = new SimpleOrderedMap<>();

      Map.Entry<String, Object>[] arr =
          (Map.Entry<String, Object>[])
              Array.newInstance(NamedList.NamedListEntry.class, rb.resultIds.size());
      // Will be set to true if there is at least one response with PURPOSE_GET_DEBUG
      boolean hasGetDebugResponses = false;

      for (ShardRequest sreq : rb.finished) {
        for (ShardResponse srsp : sreq.responses) {
          if (srsp.getException() != null) {
            // can't expect the debug content if there was an exception for this request
            // this should only happen when using shards.tolerant=true
            continue;
          }
          NamedList<Object> sdebug =
              (NamedList<Object>)
                  SolrResponseUtil.getSubsectionFromShardResponse(rb, srsp, "debug", true);

          info = (NamedList<Object>) merge(sdebug, info, EXCLUDE_SET);
          if ((sreq.purpose & ShardRequest.PURPOSE_GET_DEBUG) != 0) {
            hasGetDebugResponses = true;
            if (rb.isDebugResults() && sdebug != null) {
              NamedList<Object> sexplain = (NamedList<Object>) sdebug.get("explain");
              SolrPluginUtils.copyNamedListIntoArrayByDocPosInResponse(sexplain, rb.resultIds, arr);
            }
          }
        }
      }

      if (rb.isDebugResults()) {
        explain = SolrPluginUtils.removeNulls(arr, new SimpleOrderedMap<>());
      }

      if (!hasGetDebugResponses) {
        if (info == null) {
          info = new SimpleOrderedMap<>();
        }
        // No responses were received from shards. Show local query info.
        SolrPluginUtils.doStandardQueryDebug(
            rb.req, rb.getQueryString(), rb.wrap(rb.getQuery()), rb.isDebugQuery(), info);
        if (rb.isDebugQuery() && rb.getQparser() != null) {
          rb.getQparser().addDebugInfo(info);
        }
      }
      if (rb.isDebugResults()) {
        // put()
        int idx = info.indexOf("explain", 0);
        if (idx >= 0) {
          info.setVal(idx, explain);
        } else {
          info.add("explain", explain);
        }
      }

      rb.setDebugInfo(info);
      rb.rsp.add("debug", rb.getDebugInfo());
    }
  }

  private NamedList<String> getTrackResponse(ShardResponse shardResponse) {
    NamedList<String> namedList = new SimpleOrderedMap<>();
    if (shardResponse.getException() != null) {
      namedList.add("Exception", shardResponse.getException().getMessage());
      return namedList;
    }
    NamedList<Object> responseNL = shardResponse.getSolrResponse().getResponse();
    NamedList<?> responseHeader = (NamedList<?>) responseNL.get("responseHeader");
    if (responseHeader != null) {
      namedList.add("QTime", responseHeader.get("QTime").toString());
    }
    namedList.add("ElapsedTime", String.valueOf(shardResponse.getSolrResponse().getElapsedTime()));
    namedList.add(
        "RequestPurpose", shardResponse.getShardRequest().params.get(CommonParams.REQUEST_PURPOSE));
    SolrDocumentList docList =
        (SolrDocumentList) shardResponse.getSolrResponse().getResponse().get("response");
    if (docList != null) {
      namedList.add("NumFound", String.valueOf(docList.getNumFound()));
    }
    namedList.add("Response", String.valueOf(responseNL));
    return namedList;
  }

  @SuppressWarnings("unchecked")
  protected Object merge(Object source, Object dest, Set<String> exclude) {
    if (source == null) return dest;
    if (dest == null) {
      if (source instanceof NamedList) {
        dest = source instanceof SimpleOrderedMap ? new SimpleOrderedMap<>() : new NamedList<>();
      } else {
        return source;
      }
    } else {

      if (dest instanceof Collection) {
        // merge as Set
        if (!(dest instanceof Set)) {
          dest = new LinkedHashSet<>((Collection<?>) dest);
        }
        if (source instanceof Collection) {
          ((Collection<Object>) dest).addAll((Collection<?>) source);
        } else {
          ((Collection<Object>) dest).add(source);
        }
        return dest;
      } else if (source instanceof Number) {
        if (dest instanceof Number) {
          if (source instanceof Double || dest instanceof Double) {
            return ((Number) source).doubleValue() + ((Number) dest).doubleValue();
          }
          return ((Number) source).longValue() + ((Number) dest).longValue();
        }
        // fall through
      } else if (source instanceof String) {
        if (source.equals(dest)) {
          return dest;
        }
        // fall through
      }
    }

    if (source instanceof NamedList<?> sl && dest instanceof NamedList) {
      NamedList<Object> tmp = new NamedList<>();
      @SuppressWarnings("unchecked")
      NamedList<Object> dl = (NamedList<Object>) dest;
      // TODO simplify; drop the same-index optimization and use a Map.compute()
      for (int i = 0; i < sl.size(); i++) {
        String skey = sl.getName(i);
        if (exclude.contains(skey)) continue;
        Object sval = sl.getVal(i);
        int didx = -1;

        // optimize case where elements are in same position
        if (i < dl.size()) {
          String dkey = dl.getName(i);
          if (Objects.equals(skey, dkey)) {
            didx = i;
          }
        }

        if (didx == -1) {
          didx = dl.indexOf(skey, 0);
        }

        if (didx == -1) {
          tmp.add(skey, merge(sval, null, Collections.emptySet()));
        } else {
          dl.setVal(didx, merge(sval, dl.getVal(didx), Collections.emptySet()));
        }
      }
      dl.addAll(tmp);
      return dl;
    }

    // only add to list if JSON is different
    if (source.equals(dest)) return source;

    // merge unlike elements in a list
    List<Object> t = new ArrayList<>();
    t.add(dest);
    t.add(source);
    return t;
  }

  /////////////////////////////////////////////
  ///  SolrInfoBean
  ////////////////////////////////////////////

  @Override
  public String getDescription() {
    return "Debug Information";
  }

  @Override
  public Category getCategory() {
    return Category.OTHER;
  }
}
